import paddle
from paddle import nn

class DecoderRNN(nn.Layer):
    def __init__(self, embed_size, hidden_size, vocab_size, num_layers, max_seq_length=20):
        super(DecoderRNN, self).__init__();
        self.embed = nn.Embedding(vocab_size, embed_size)
        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, time_major=False)
        self.linear = nn.Linear(hidden_size, vocab_size)
        self.max_seq_length = max_seq_length

    def forward(self, features, captions, lengths):
        """Decode image feature vectors and generates captions."""
        embeddings = self.embed(captions)
        embeddings = paddle.concat((paddle.unsqueeze(features,1), embeddings), 1)
        hiddens, _ = self.lstm(embeddings, sequence_length=lengths)
        outputs = self.linear(hiddens[0])
        return outputs

    def sample(self, features, states=None):
        """Generage captions for given image features using greedy search."""
        sampled_ids = []
        inputs = paddle.unsqueeze(features, 1)
        for i in range(self.max_seq_length):
            hiddens, states = self.lstm(inputs, states)
            outputs = self.linear(paddle.unsqueeze(hiddens, 1))
            _, predicted = paddle.max(outputs, 1)
            sampled_ids.append(predicted)
            inputs = self.embed(predicted)
            inputs = paddle.unsqueeze(inputs, 1)
        sampled_ids = paddle.stack(sampled_ids, 1)
        return sampled_ids

